#include <torch/script.h>  // One-stop header.
#include <iostream>
#include <memory>
#include <chrono>


using namespace std::chrono;

void testDemo();

int main() {
  // // Deserialize the ScriptModule from a file using torch::jit::load().
  // //std::shared_ptr<torch::jit::script::Module> module =
  // //          torch::jit::load("E:/HM_DL/torch_test/traced_resnet_model.pt");

  // using torch::jit::script::Module;
  // Module module =
  //     torch::jit::load("/home/coshe/PycharmProjects/LSTM_Test/model/best_model.pt");

  // std::cout << "ok\n";
  // // Create a vector of inputs.
  // std::vector<torch::jit::IValue> inputs;
  // inputs.push_back(torch::rand({1, 10, 1}));

  // // Execute the model and turn its output into a tensor.
  // std::cout << "inputs: " << std::endl;
  // std::cout << inputs << std::endl;

  // at::Tensor output = module.forward(inputs).toTensor();



  // std::cout << output << '\n';
  // // while (1)
  // //   ;



  testDemo();


}


void testDemo(){
  const float PI = 3.1415926;
  using torch::jit::script::Module;
  torch::jit::script::Module module =
      torch::jit::load("/home/coshe/PycharmProjects/LSTM_Test/model/best_model.pt");
  
  
  std::vector<float> input_data;

  //for (int i = 0; i < 100; i++){
    for (int j = 0; j < 10; j ++){
      input_data.push_back(std::sin(PI/9.0 * j));
    }
  //}

  std::cout << input_data << std::endl;

  std::cout << std::endl;
  std::cout << std::endl;  
  std::cout << std::endl;
  std::cout << std::endl;

  std::cout << "-------------------------------------------------" << std::endl;

  at::Tensor inputTensor = torch::from_blob(input_data.data(), {1, 10});

  auto start = system_clock::now();
  
  at::Tensor output = module.forward({inputTensor}).toTensor();

  auto end = system_clock::now();

  auto dur = end - start;

  std::chrono::milliseconds ms = std::chrono::duration_cast<std::chrono::milliseconds>(dur);

  std::cout << "Time: " << ms.count() << "ms" << std::endl;

  std::cout << "pred: " << output << std::endl;
  std::cout << "real: " << std::sin(PI/9.0 * 10) << std::endl;


}






